from torch.utils.data import DataLoader
from torchvision import models
from torch import optim
from torch.autograd import Variable
import torch.nn as nn
from tqdm import tqdm
import os.path as osp
from dataset import UserDataset, userDataset
from net import FCN8s
from globalData import lr, class_num, checkpoint_folder, DEVICE
from utils import save_checkpoint, load_checkpoint


class AverageMeters(object):
    def __init__(self):
        self.reset()

    def reset(self):
        self.avg = 0
        self.num = 0
        self.loss = 0

    def update(self, value):
        self.loss += value
        self.num += 1
        self.avg = round(self.loss / self.num, 2)
        return self.avg


class TrainUtils:
    def __init__(self, userDataset: UserDataset, batchSize: int, fcn: FCN8s):
        self.dataLoader = DataLoader(userDataset, batch_size=batchSize, num_workers=0)
        self.device = DEVICE
        self.fcn = fcn.to(self.device)
        self.optimizer = optim.Adam(fcn.parameters(), lr=1e-4)
        self.lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer, mode='min', factor=0.1, patience=5, verbose=True,
            threshold=0.002, threshold_mode='rel'
        )
        self.calLoss = nn.CrossEntropyLoss().to(self.device)  # .to(self.device) <====> .cuda()

    def trainNet(self):
        for epoch in range(30):
            log_values = AverageMeters()
            # if epoch % 10 == 0 and epoch != 0:
            #     for group in self.optimizer.param_groups:
            #         group['lr'] *= 0.5
            pbar = tqdm(self.dataLoader, 'Epoch ' + str(epoch), ncols=80)
            self.fcn.train()  # 设置为训练过程。此时网络 的 参数能够进行参数学习
            for i, sample in enumerate(pbar):
                #  Variable为tensor数据构建计算图，便于网络的运算 https://www.cnblogs.com/czz0508/p/10333359.html
                imgdata, imglabel = Variable(sample['img']).to(self.device), Variable(sample['label'].long()).to(
                    self.device)
                self.optimizer.zero_grad()  # 梯度清0
                netOutput = self.fcn(imgdata)  # 将训练图片 送入网络
                loss = self.calLoss(netOutput, imglabel)  # 计算损失, label 和 网络输出的预测的结果进行对比，计算loss
                log_values.update(loss.item())
                loss.backward()  # 误差反向传播
                pbar.set_postfix(
                    {'train loss': f'{log_values.avg:.4f}'}, refresh=False)
                self.optimizer.step()  # 用优化器去更新权重参数
            self.lr_scheduler.step(log_values.avg)
            save_checkpoint(epoch, self.fcn, self.optimizer, self.lr_scheduler, checkpoint_folder)
        print('*' * 10, 'train  end')


if __name__ == '__main__':
    # userDataset 为上一个部分，得到的数据集对象
    pretrained_net = models.vgg16_bn(pretrained=False)
    batchSize = 4
    net = FCN8s(pretrained_net, class_num)
    trainUtils = TrainUtils(userDataset, batchSize, net)
    trainUtils.trainNet()
